from openai import OpenAI
import pandas as pd
import re

def predbailian(prompt):
    bailian_api_key = 'sk-fc0b61726c1d45b2a11946e8caa5ec62'
    client = OpenAI(
        api_key=bailian_api_key,
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    )
    completion = client.chat.completions.create(
        model="qwen-plus",
        messages=[
            {'role': 'system', 'content': 'You are a helpful assistant.'},
            {'role': 'user', 'content': prompt}],
        temperature=0.001
    )
    content = completion.choices[0].message.content
    return content
class Predbailian():
    def __init__(self):
        pass
    def cleanres(self,data, res):
        splitresid = [r.split('--')[0].replace(' ','') for r in res.split('\n') if re.findall('\d', r)]
        splitres = [r.split('--')[1].replace(' ','') for r in res.split('\n') if re.findall('\d', r)]
        resdf = pd.DataFrame({'id': splitresid, 'res': splitres})
        resdf['id'] = (resdf['id'].astype(int)/2).astype(int)
        lastres = data.merge(resdf, how='left', on='id')
        return lastres

    def groupres(self,listres):
        group = 0
        group_list = []
        for i in range(len(listres)):
            if i == 0:
                group_list.append(group)
            elif (listres[i] == '问询') and (listres[i-1] == '问询') :
                group_list.append(group)
            elif (listres[i] == '问询') and (listres[i-1] == '应答'):
                group += 1
                group_list.append(group)
            elif (listres[i] == '应答') and (listres[i-1] == '问询'):
                group_list.append(group)
            elif (listres[i] == '应答') and (listres[i-1] == '应答'):
                group += 1
                group_list.append(group)

        return group_list



    def distinguishingTopicGroups(self,data):
        print(data)
        print(type(data))
        df = pd.DataFrame(data)
        df['id'] = df.index + 1
        dia, dialog = [], []
        for i, row in df.iterrows():
            dia.append(' --' + ' 员工: ' + row['learnerDialogue'])
            dia.append(' --' + ' 客户: ' + row['robotTalk'])

        for i in range(len(dia)):
            dialog.append(str(i+1) + dia[i])


        dialogtext = '\n'.join(dialog)
        prompt = f'''
        以下是一段客户和员工的聊天内容，格式为
            id -- 员工说话的内容
            id -- 客户说话的内容 
        任务：请判断客户说话的内容是哪一种类型，属于应答，问询两种类型的哪一种。
    
        问询定义：客户向员工获取信息或者寻求帮助，例如：可以帮忙安排一下吗？ 
        应答定义：客户没有对员工提出问题，例如：非常感谢，我没有需要了解的信息了，谢谢！你好，我姓李。
        
        只要客户说话的内容中出现问询相关的内容 例如有问号， 类型标注为 问询
        请按照给定的输出格式输出结果
        请针对每一行的客户说话的内容进行类型标注，请不要遗漏任何结果。
    
        --------------------------------------------
        聊天内容:
        {dialogtext}
        --------------------------------------------
    
        输出格式示例：
            id -- 应答
            id -- 应答
        '''

        qwenresponse = predbailian(prompt)
        print(qwenresponse)

        res = self.cleanres(df, qwenresponse)
        res['groupid'] = self.groupres(res['res'])
        res['Groupcount'] = res.groupby('groupid')['nodeDetailId'].transform('count')
        res.loc[res['Groupcount'] != 1, 'seq'] = res[res['Groupcount'] != 1].groupby('groupid').cumcount() + 1
        letters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
        groupids = res[res['Groupcount'] != 1]['groupid'].unique()
        mapping = dict(zip(sorted(groupids), letters[:len(groupids)]))
        res.loc[res['Groupcount'] != 1, 'Groupid'] = res[res['Groupcount'] != 1]['groupid'].map(mapping)
        res['GroupName'] = res['Groupid'] + res['seq'].astype(str).str.strip('.0')
        res['GroupName'] = res['GroupName'].fillna('')

        return res[['nodeDetailId', 'GroupName']].to_dict(orient='records')


if __name__ == '__main__':
    Predbailian()
